import os, sys, commands, ntpath, math, hashlib, unicodedata

__version__ = '1.0.0'
__author__ = "joesox@gmail.com"
__url__ = 'www.joeswammi.com'
"""
writing on Python 2.6.5
PREREQUISITES:
TrID (and its .trd file)- File Identifier: http://mark0.net/soft-trid-e.html
ssdeep: http://code.google.com/p/pyssdeep/
Pyrex: http://www.cosc.canterbury.ac.nz/greg.ewing/python/Pyrex/
"""


class fileinspecttools:
    """
    based from 'Building a Malware Zoo' whitepaper
    http://www.sans.org/reading_room/whitepapers/malicious/building-malware-zoo_33543
    """
    def __repr__(self):
        if not self:  
            return 'Attrs()'  
        return 'Attrs(%s)' % list.__repr__(self) 

    def __init__(self):
        self.value = ""
        self.FS_ROOT = 'c:\\'
        self.TRIDPATH = os.path.join(os.getcwd(), "trid")
        self.SSDEEPPATH = os.path.join(os.getcwd(), "ssdeep", "ssdeep")
        self.DATABASEFILE = os.path.join(os.getcwd(), "ssdeep","database.dat")
    
    def getSize(self, filename):
        #Return file size
        size = os.path.getsize(filename)
        return size
    
    def getFileType(self,filename):
        #Execute "file -b <filename>"
        cmd = self.TRIDPATH + " -d: "+filename
        pipe = os.popen(cmd)
        results = pipe.read()
        return results
    
    def getHash(self, filename):
        data=open(filename, "r").read()
        #Compute MD5 and SHA1 hashes
        md5hash=hashlib.md5(data).hexdigest()
        sha1hash=hashlib.sha1(data).hexdigest()
        #Build return structure
        results=[md5hash,sha1hash]
        return results
    
    def getFuzzyHash(self, filename):
        match     = ""
        #Compute fuzzy hashes and fuzzy matches
        cmd = self.SSDEEPPATH + " -bm "+ self.DATABASEFILE + " " + filename
        if(os.path.isfile(self.DATABASEFILE)):
            pipe = os.popen(cmd)
            match = pipe.read()
        else:
            #create the file with the header
            fo = open(self.DATABASEFILE, "wb")
            fo.write("ssdeep,1.1--blocksize:hash:hash,filename\n")
            fo.close()
            
        cmd = self.SSDEEPPATH + " -b " + filename
        pipe = os.popen(cmd)
        fuzzyHash = pipe.read().split("\n")
        
        #add to db
        # Open a file
        if(match == ""):
            fo = open(self.DATABASEFILE, "a+")
            fo.write(fuzzyHash[1] + "\n")
            fo.close()

        results = [fuzzyHash[1], match]
        return results
    
    def extractStrings(self,filename):
        frag = ""
        strList = []
        bufLen = 2048
        FRAG_LEN = 4 #Min length to report as string
        fp = open(filename, "rb")
        offset = 0
        buf = fp.read(bufLen)
        while buf:
            for char in buf:
                #Uses curses library to locate printable chars in binary files
                try:
                    if (unicodedata.category(unicode(char))[0] in 'LNPS')==False:
                        if len(frag)>FRAG_LEN:
                            strList.append([hex(offset-len(frag)),frag])
                            frag = ""
                    else:
                        frag = frag + char
                        offset+=1
                        buf=fp.read(bufLen)
                except ValueError:
                    pass # UnicodeDecodeError: 'ascii' codec can't decode byte

        return strList
    
    def getEntropy(self,filename):
        """ 
        Calculate the entropy of a chunk of data. 
        Generally, the higher the entropy the greater 
        the chance the sample is compressed or encrypted
        """
        entropy = 0
        
        frag = ""
        strList = []
        bufLen = 2048
        FRAG_LEN = 4 #Min length to report as string
        fp = open(filename, "rb")
        offset = 0
        buf = fp.read(bufLen)

        #Uses curses library to locate printable chars in binary files
        try:
            for x in range(256):
                p_x = float(buf.count(chr(x)))/len(buf)
                if p_x > 0:
                    entropy += - p_x*math.log(p_x,2)
        except ValueError:
            pass # UnicodeDecodeError: 'ascii' codec can't decode byte  
        
        fp.close()
        return entropy
    
    def Examine(self, filename):
        print "Analysing " + filename
        print "  Filesize: " + str(self.getSize(filename)) + " bytes"
        print "  Filetype: " + str(self.getFileType(filename))
        hashes = self.getHash(filename)
        print "  MD5 hash: " + str(hashes[0])
        print "  SHA1 hash: " + str(hashes[1])
        fuzzyhashes= self.getFuzzyHash(filename)
        print "  Fuzzy hash: " + str(fuzzyhashes[0])
        print "   Previous match: " + str(fuzzyhashes[1])
        print "  Strings found: " + str(self.extractStrings(filename))
        print "  Entropy [higher=compressed or encrypt]: " + str(self.getEntropy(filename))
        
def test():
    t=fileinspecttools()
    print "*******"
    t.Examine(os.path.join(os.getcwd(), "samples", "sample.bmp"))
    print "*******"
    t.Examine(os.path.join(os.getcwd(), "samples", "sample.7z"))
    #t.Examine(os.path.join(os.getcwd(), "trid.exe"))
    
    print "*******"
if __name__ == '__main__':
    test()
